#pragma once

#define MAX_DPA_THREAD_NUM 256
#define DEFAULT_BATCH_SIZE 1

#include <stdlib.h>
#include <gflags/gflags.h>

#include <doca_log.h>
#include <infiniband/verbs.h>
#include <infiniband/mlx5_api.h>
#include <libflexio/flexio.h>
#include <numa.h>

#include "DeviceShare.h"
#include "DpaCommon.h"
#include "DocaSdkCommon.h"
#include "glog/logging.h"

extern "C"
{
    flexio_func_t FlexioProcess;
    flexio_func_t thread_task_init;
    flexio_func_t get_max_thread_time;
    doca_dpa_func_t DocaSdkProcess;
}
extern struct flexio_app * mt_thread_device;

DECLARE_string(rdma_device);

class DpaUnit{
public:
    int get_thread_num() {return active_thread_num;}
    virtual void register_host_memory(void * ptr, size_t size) = 0;
    virtual void submit_task(uint64_t ddata, int execute_thread_num = 1) = 0;
    virtual void init_dpa_thread_info(DpaThreadInfo & T) = 0;
    virtual void poll_thread() = 0;
    virtual size_t get_execute_cycle() = 0;

public:
    int active_thread_num;
};

class DpaFlexioUnit : public DpaUnit {
public:
    DpaFlexioUnit() {
        this->active_thread_num = MAX_DPA_THREAD_NUM;
        this->device_name = FLAGS_rdma_device;
        app = mt_thread_device;
        func = FlexioProcess;
        init_func = thread_task_init;

        this->context_init(this->app);
        this->log_init(0);
        this->command_queue_init();
    }

    void context_init(flexio_app *app_device) {
        ctx = new FLEX::Context(device_name);
        ctx->alloc_pd();
        ctx->create_process(app_device);
        ctx->create_window();
        ctx->generate_flexio_uar();
    }

    void log_init(int num) { //给多少个函数进行log输出
        ctx->print_init("mt_memory", num);
    }

    void command_queue_init() {
        cmdq = new FLEX::CommandQ(ctx->get_process(), active_thread_num, DEFAULT_BATCH_SIZE);
    }

    /* malloc在dma进行，这里只需要注册，映射好 */
    void register_host_memory(void * ptr, size_t size) override{
        ctx->reg_mr(ptr, (size + 63) & (~63));
        LOG(INFO) << "successful finish register host memory";
    }

    /* 实验证明分配多个lkey不成功，还是只能用一大片内存，同一个window */
    uint32_t get_lkey() {
        return ctx->get_mr(0)->lkey;
    }

    uint32_t get_window_id() {
        return ctx->get_window_id();
    }

    void submit_task(uint64_t ddata, int execute_thread_num) override{
        (void) ddata;
        for (int i = 0; i < execute_thread_num; ++i) {
            cmdq->add_task(func, i);
        }
    }

    void init_dpa_thread_info(DpaThreadInfo & T) override {
        T.window_id       = this->get_window_id();
        T.input_mem_mkey  = this->get_lkey();
        T.output_mem_mkey = this->get_lkey();

        uint64_t rpc_ret_val;
        flexio_uintptr_t dev_config_data = ctx->move_to_dev(T);
        assert(flexio_process_call(ctx->get_process(), init_func, &rpc_ret_val,
            dev_config_data) == FLEXIO_STATUS_SUCCESS);
    }

    void poll_thread() {
        cmdq->run();
        auto time_us = cmdq->wait_run(300);
        (void)time_us;
    }

    size_t get_execute_cycle() override {
        size_t execute_cycle = 0;
        assert(flexio_process_call(ctx->get_process(), get_max_thread_time, &execute_cycle) == FLEXIO_STATUS_SUCCESS);
        return execute_cycle;
    }

    ~DpaFlexioUnit() {
        
    }


public:
    flexio_func_t * func;
    flexio_func_t * init_func;
    struct flexio_app * app;
    std::string device_name;
    FLEX::Context * ctx;
    FLEX::CommandQ * cmdq;
};

class DpaDocasdkUnit : public DpaUnit{
public:
    DpaDocasdkUnit() {
        this->active_thread_num = 128; //TODO: max_thread_num_per_kernel = 128
        	
        doca_error_t result;

        strcpy(cfg.pf_device_name, DEVICE_DEFAULT_NAME);
        strcpy(cfg.rdma_device_name, DEVICE_DEFAULT_NAME);

        result = doca_log_backend_create_standard();
        result = doca_log_backend_create_with_file_sdk(stderr, &sdk_log);
        result = doca_log_backend_set_sdk_level(sdk_log, DOCA_LOG_LEVEL_WARNING);

        result = allocate_dpa_resources(&cfg, &resources);
        if (result != DOCA_SUCCESS) {
            LOG(INFO) << "Failed to Allocate DPA Resources";
        }

        LOG(INFO) << "finish dpa allocate";

        unsigned int value = 0;
        doca_dpa_get_max_threads_per_kernel(resources.pf_dpa_ctx, &value);
        LOG(INFO) << "max thread per kernel = " << value;

        unsigned long long max_time = 0;
        doca_dpa_get_kernel_max_run_time(resources.pf_dpa_ctx, &max_time);
        LOG(INFO) << "max execute time per kernel = " << max_time;

        comp_event = NULL;
        comp_event_val = 10;

        result = create_doca_dpa_completion_sync_event(resources.pf_dpa_ctx,
                                resources.pf_doca_device,
                                &comp_event,
                                NULL);
        if (result != DOCA_SUCCESS) {
            LOG(INFO) << "sync_event_completion failed";
        }

        func = DocaSdkProcess;
        init_func = thread_task_init;

        LOG(INFO) << "finish sync event register";
    }

    void register_host_memory(void * ptr, size_t size) override {
        void * send_ptr = ptr;
        host_mmap_obj.mmap_type = MMAP_TYPE_CPU;
        host_mmap_obj.doca_dpa = resources.pf_dpa_ctx;
        host_mmap_obj.doca_device = resources.pf_doca_device;
        host_mmap_obj.permissions = DOCA_ACCESS_FLAG_LOCAL_READ_WRITE | DOCA_ACCESS_FLAG_RDMA_WRITE |
                        DOCA_ACCESS_FLAG_RDMA_READ | DOCA_ACCESS_FLAG_RDMA_ATOMIC | DOCA_ACCESS_FLAG_PCI_READ_WRITE;
        host_mmap_obj.memrange_addr = send_ptr;
        host_mmap_obj.memrange_len = ( (size + 63) & (~63) );

        doca_error_t cur_doca_err = doca_mmap_obj_init(&host_mmap_obj);
        if (cur_doca_err != DOCA_SUCCESS) {printf ("register_host_memory error!\n"); fflush(stdout);}
    }

    void submit_task(uint64_t ddata, int execute_thread_num) override {
        (void) ddata;
        doca_error_t result;
        result = doca_dpa_kernel_launch_update_set(resources.pf_dpa_ctx,
						   NULL,
						   0,
						   comp_event,
						   comp_event_val,
						   execute_thread_num,
						   func,
						   host_mmap_obj.dpa_mmap_handle
						   );
        if (result != DOCA_SUCCESS) {
            LOG(INFO) << "kernel_launch failed";
        }
    }

    void init_dpa_thread_info(DpaThreadInfo & T) override {
        uint64_t rpc_ret_val;
        doca_dpa_dev_uintptr_t dev_addr; 
        doca_error_t doca_err = DOCA_SUCCESS;
        doca_err = doca_dpa_mem_alloc(resources.pf_dpa_ctx, sizeof(T), &dev_addr);
        doca_dpa_h2d_memcpy(resources.pf_dpa_ctx, dev_addr, &T, sizeof(T) );
        if (doca_err != DOCA_SUCCESS) {
            LOG(INFO) << "malloc failed";
        }

        doca_err = doca_dpa_rpc(resources.pf_dpa_ctx,
				init_func,
				&rpc_ret_val,
				dev_addr);
        if (doca_err != DOCA_SUCCESS) {
            LOG(INFO) << "Function init_func failed";
        }
    }

    void poll_thread() override {
        doca_error_t result;
        result = doca_sync_event_wait_gt(comp_event, comp_event_val - 1, SYNC_EVENT_MASK_FFS);
        if (result != DOCA_SUCCESS) {
            LOG(INFO) << ("poll_wait error");
        }
    }

    size_t get_execute_cycle() override {
        uint64_t execute_cycle = 0; 
        doca_error_t doca_err = DOCA_SUCCESS;

        doca_err = doca_dpa_rpc(resources.pf_dpa_ctx,
				get_max_thread_time,
				&execute_cycle);
        if (doca_err != DOCA_SUCCESS) {
            LOG(INFO) << "Function get_max_thread_time failed";
        }
        return execute_cycle;
    }

public:
    struct dpa_config cfg;
    struct dpa_resources resources;
    struct doca_mmap_obj host_mmap_obj;
    struct doca_sync_event * comp_event;
    struct doca_log_backend * sdk_log;
    uint64_t comp_event_val;

    doca_dpa_func_t * func;
    doca_dpa_func_t * init_func;
};